from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QWidget, QMainWindow, QLabel, QFileDialog, QVBoxLayout, QApplication, QSlider
from PyQt5.QtGui import QPainter, QPixmap, QImage
from PyQt5.QtCore import Qt, QRect, qDebug
import cv2 as cv
import numpy as np
import vtkmodules.all as vtk
from vtkmodules.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor

class MyLabel(QLabel):
    def __init__(self, parent=None):
        super(MyLabel, self).__init__(parent)
        self.srcImage = None
        self.bTempPixmap = None
        self.ZoomValue = 1.0
        self.XPtInterval = 0
        self.YPtInterval = 0
        self.Pressed = False
        self.bZoom = True
        self.OldPos = None
        self.nMovePos = None
        self.intPenWidth = 1



    def SetPic(self, image):
        self.srcImage = image
        self.update()

    def setZoom(self, bZoom):
        self.bZoom = bZoom

    def paintEvent(self, event):
        painter = QPainter(self)
        tmpPixMap = QPixmap(self.width(), self.height())
        tmpPixMap.fill(Qt.transparent)
        painter.begin(tmpPixMap)

        if self.srcImage != None:
            width = min(self.srcImage.width(), self.width())
            height = width * 1.0 / (self.srcImage.width() * 1.0 / self.srcImage.height())
            height = min(height, self.height())
            width = height * 1.0 * (self.srcImage.width() * 1.0 / self.srcImage.height())

            painter.translate(self.width() / 2 + self.XPtInterval, self.height() / 2 + self.YPtInterval)
            #painter.translate(width, height)
            painter.scale(self.ZoomValue, self.ZoomValue)

            picRect = QRect(int(-width / 2), int(-height / 2), int(width), int(height))

            painter.drawImage(picRect, self.srcImage)
        painter.end()



    def wheelEvent(self, event):
        if self.bZoom == False:
            return
        value = event.angleDelta().y()

        if value > 0:
            self.OnZoomInImage()
        else:
            self.OnZoomOutImage()

        self.update()

    def mouseDoubleClickEvent(self, event):
        self.OnPresetImage()

    def mousePressEvent(self, event):
        self.OldPos = event.pos()
        self.Pressed = True

    def mouseMoveEvent(self, event):
        self.setCursor(Qt.SizeAllCursor)
        pos = event.pos()
        xPtInterval = pos.x() - self.OldPos.x()
        yPtInterval = pos.y() - self.OldPos.y()

        self.XPtInterval += xPtInterval
        self.YPtInterval += yPtInterval

        self.OldPos = pos
        self.update()

    def mouseReleaseEvent(self, event):
        self.Pressed = False
        self.setCursor(Qt.ArrowCursor)


    def OnZoomInImage(self):
        self.ZoomValue += 0.05
        self.update()

    def OnZoomOutImage(self):
        self.ZoomValue -= 0.05
        if self.ZoomValue <= 0:
            self.ZoomValue = 0.05
            return

        self.update()

    def OnPresetImage(self):
        self.ZoomValue = 1.0
        self.XPtInterval = 0
        self.YPtInterval = 0
        self.update()



class VtkPointCloud:

    def __init__(self):
        self.vtkPolyData = vtk.vtkPolyData()
        self.clearPoints()
        # mapper = vtk.vtkPolyDataMapper()
        # mapper.SetInputData(self.vtkPolyData)
        # mapper.SetColorModeToDefault()
        # mapper.SetScalarRange(zMin, zMax)
        # mapper.SetScalarVisibility(1)
        self.vtkActor = vtk.vtkActor()
        # self.vtkActor.GetProperty().SetPointSize(4)
        # self.vtkActor.SetMapper(mapper)

    def setRange(self, zMin=-0.2, zMax=1):
        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputData(self.vtkPolyData)
        mapper.SetColorModeToDefault()
        mapper.SetScalarRange(zMin, zMax)
        mapper.SetScalarVisibility(1)
        self.vtkActor.GetProperty().SetPointSize(4)
        self.vtkActor.SetMapper(mapper)

    def addPoint(self, point):
        pointId = self.vtkPoints.InsertNextPoint(point[:])
        self.vtkDepth.InsertNextValue(point[2])
        self.vtkCells.InsertNextCell(1)
        self.vtkCells.InsertCellPoint(pointId)

        self.vtkCells.Modified()
        self.vtkPoints.Modified()
        self.vtkDepth.Modified()

    def clearPoints(self):
        self.vtkPoints = vtk.vtkPoints()
        self.vtkCells = vtk.vtkCellArray()
        self.vtkDepth = vtk.vtkDoubleArray()
        self.vtkDepth.SetName('DepthArray')
        self.vtkPolyData.SetPoints(self.vtkPoints)
        self.vtkPolyData.SetVerts(self.vtkCells)
        self.vtkPolyData.GetPointData().SetScalars(self.vtkDepth)
        self.vtkPolyData.GetPointData().SetActiveScalars('DepthArray')

class VTKWindow(QMainWindow):
    def setupUi(self, MainWindow):
        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(1600, 800)
        self.centralWidget = QtWidgets.QWidget(MainWindow)
        self.centralWidget.setObjectName("centralWidget")

        self.pushButton_tif = QtWidgets.QPushButton(self.centralWidget)
        self.pushButton_tif.setGeometry(QtCore.QRect(15, 15, 120, 30))
        self.pushButton_tif.setMinimumSize(QtCore.QSize(0, 23))
        self.pushButton_tif.setAutoRepeatInterval(100)
        self.pushButton_tif.setObjectName("pushButton_tif")


        self.widgetDisplay = QtWidgets.QWidget(self.centralWidget)
        # self.widgetDisplay.setStyleSheet('''QWidget{background-color:rgb(0, 0, 0);}''')
        self.widgetDisplay.setGeometry(QtCore.QRect(850, 50, 700, 600))
        self.widgetDisplay.setObjectName("widgetDisplay")

        # self.label = QtWidgets.QLabel(self.centralWidget)
        # self.label.setGeometry(QtCore.QRect(50, 700, 1000, 50))
        font = QtGui.QFont()
        font.setFamily("黑体")
        font.setPointSize(12)
        # self.label.setFont(font)
        # self.label.setText("")
        # self.label.setObjectName("label")

        self.widgetImage = QtWidgets.QWidget(self.centralWidget)
        self.widgetImage.setGeometry(QtCore.QRect(50, 50, 700, 600))
        self.widgetImage.setObjectName("widgetImage")

        vbox = QVBoxLayout()


        self.dstImage = np.array([])
        self.grayImage = np.array([])
        self.rgbImage = np.array([])

        self.slider = QSlider(Qt.Horizontal)
        self.slider.setMinimum(0)
        # 设置最大值
        self.slider.setMaximum(100)
        # 步长
        self.slider.setSingleStep(5)
        # 设置当前值
        self.slider.setValue(50)

        # 刻度位置，刻度在下方
        self.slider.setTickPosition(QSlider.TicksBelow)
        # 设置刻度间隔
        self.slider.setTickInterval(5)

        vbox.addWidget(self.slider)
        self.slider.valueChanged.connect(self.updateImage)

        self.imageLabel = MyLabel()
        # self.imageLabel.setGeometry(QtCore.QRect(50, 50, 700, 600))
        self.imageLabel.setObjectName("imageLabel")
        # self.imageLabel.setParent(self.imageLabel)
        self.imageLabel.setStyleSheet('''QLabel{background-color:rgb(50, 80, 100);}''')
        vbox.addWidget(self.imageLabel)
        self.widgetImage.setLayout(vbox)


        MainWindow.setCentralWidget(self.centralWidget)


        self.vtk_show = QVTKRenderWindowInteractor(self.widgetDisplay)
        self.vtk_show.resize(self.widgetDisplay.size())

        self.pointCloud = VtkPointCloud()
        self.mapper = None
        self.actor = None
        self.showPoints = None
        self.edgeActor = None
        self.renderer = None
        self.i_ren = None
        self.filename = ""

        self.colorbg = (0, 0, 0)

        self.pushButton_tif.clicked.connect(self.read_tif)


        self.setMouseTracking(True)

        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

        self.init()

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "MainWindow"))
        self.pushButton_tif.setText(_translate("MainWindow", "打开tif文件"))


    def init(self):
        self.renderer = vtk.vtkRenderer()
        self.renderer.AddActor(self.pointCloud.vtkActor)
        self.renderer.SetBackground(.2, .3, .4)
        self.renderer.ResetCamera()

        # Render Window
        self.vtk_show.SetInteractorStyle(vtk.vtkInteractorStyleMultiTouchCamera())
        self.vtk_show.GetRenderWindow().AddRenderer(self.renderer)
        self.i_ren = self.vtk_show.GetRenderWindow().GetInteractor()
        self.vtk_show.Render()
        self.i_ren.Initialize()
        self.vtk_show.Start()

    def cvToQImage(self, data):
        # 8-bits unsigned, NO. OF CHANNELS=1
        if data.dtype == np.uint8:
            channels = 1 if len(data.shape) == 2 else data.shape[2]
        if channels == 3:  # CV_8UC3
            # Copy input Mat
            # Create QImage with same dimensions as input Mat
            img = QImage(data, data.shape[1], data.shape[0], data.strides[0], QImage.Format_RGB888)
            return img.rgbSwapped()
        elif channels == 1:
            # Copy input Mat
            # Create QImage with same dimensions as input Mat
            data = cv.cvtColor(data, cv.COLOR_BGR2GRAY)
            img = QImage(data, data.shape[1], data.shape[0], data.shape[1], QImage.Format_Indexed8)
            # print(data)
            return img
        else:
            qDebug("ERROR: numpy.ndarray could not be converted to QImage. Channels = %d" % data.shape[2])
            return QImage()


    def depthToColor(self, depthImg, min, max):
        x = depthImg

        # newArray = np.where(x < -30000, min, x)
        # minIndex = np.argwhere(x < -30000)

        img_dep = (depthImg - min) / (max - min) * 255
        img_dep = img_dep.astype(np.uint8)
        result = cv.applyColorMap(img_dep, cv.COLORMAP_JET)
        result[np.where(x < -30000)] = 0

        return result




    def read_tif(self):
        self.read_file("jpg")
        if self.filename != "":
            self.grayImage = cv.imdecode(np.fromfile(self.filename, dtype=np.uint8), cv.IMREAD_COLOR)
            depthFilename = self.filename.replace("GA", "HA")
            depthFilename = depthFilename.replace(".jpg", ".tif")
            depthImg = cv.imdecode(np.fromfile(depthFilename, dtype=np.int16), -1)

            depthImg_u = np.unique(depthImg)

            self.rgbImage = self.depthToColor(depthImg, depthImg_u[1], depthImg_u[-1])

            self.pointCloud.clearPoints()
            self.pointCloud.setRange(-depthImg_u[-1] / 1000 + 0.05, -depthImg_u[1] / 1000 - 0.05)
            for h in range(0, depthImg.shape[0], 4):
                for w in range(0, depthImg.shape[1], 4):
                    if depthImg[h, w] > -10000:
                        point = (w * 0.01, h * 0.01, float(-depthImg[h, w] / 1000.0))
                        self.pointCloud.addPoint(point)

                        # poins.InsertNextPoint(w*0.01, h*0.01, float(depthImg[h, w] / 1000.0))

            self.dstImage = cv.addWeighted(self.rgbImage, float(self.slider.value() / 100), self.grayImage, float(100 - self.slider.value()) / 100, 0.2)

            self.imageLabel.SetPic(self.cvToQImage(self.dstImage))

            QApplication.processEvents()
            self.renderer.ResetCamera()

    def updateImage(self):

        if self.rgbImage.any() and self.grayImage.any():
            self.dstImage = cv.addWeighted(self.rgbImage, float(self.slider.value() / 100), self.grayImage, float(100 - self.slider.value()) / 100, 0.2)
            self.imageLabel.SetPic(self.cvToQImage(self.dstImage))
            QApplication.processEvents()
















    def showVTK(self):

        self.mapper = vtk.vtkPolyDataMapper()
        self.mapper.SetInputConnection(self.showPoints.GetOutputPort())
        # self.mapper.SetInputData(self.showPoints)
        self.mapper.SetColorModeToDefault()
        self.mapper.SetScalarRange(-1, 1)
        self.mapper.SetScalarVisibility(1)
        self.mapper.Update()

        self.actor = vtk.vtkActor()
        self.actor.SetMapper(self.mapper)

        self.actor.GetProperty().SetColor(1, 0, 0)
        self.actor.GetProperty().SetPointSize(2)

        self.renderer = vtk.vtkRenderer()
        self.renderer.AddActor(self.actor)
        self.renderer.SetBackground(self.colorbg)

        # 5.windows

        self.vtk_show.SetInteractorStyle(vtk.vtkInteractorStyleMultiTouchCamera())

        self.vtk_show.GetRenderWindow().AddRenderer(self.renderer)
        self.i_ren = self.vtk_show.GetRenderWindow().GetInteractor()
        self.vtk_show.Render()
        self.i_ren.Initialize()
        self.vtk_show.Start()





    def read_file(self, file_type):
        dialog = QFileDialog()
        dialog.setFileMode(QFileDialog.ExistingFile)
        dialog.setAcceptMode(QFileDialog.AcceptOpen)
        self.filename, _ = dialog.getOpenFileName(dialog, "Open file", "", "files (*." + file_type + ")")
        # self.label.setText(self.filename)



if __name__ == "__main__":
    import sys

    app = QtWidgets.QApplication(sys.argv)
    MainWindow = QtWidgets.QMainWindow()
    ui = VTKWindow()
    ui.setupUi(MainWindow)
    MainWindow.show()

    sys.exit(app.exec_())
